#include "stdafx.h"

#include "AutoLock.h"
#include "ProcessProtectDriver.h"

Globals g_Data;

VOID
ProcessProtectUnload(_In_ PDRIVER_OBJECT DriverObject);

NTSTATUS
ProcessProtectDeviceControl(_In_ PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
ProcessProtectCreateClose(_In_ PDEVICE_OBJECT DeviceObject, PIRP Irp);

OB_PREOP_CALLBACK_STATUS
OnPreOPenProcess(_In_ PVOID RegistrationContext, _Inout_ POB_PRE_OPERATION_INFORMATION OperationInformation);

BOOLEAN
AddProcess(_In_ ULONG Pid);

BOOLEAN
RemoveProcess(_In_ ULONG Pid);

BOOLEAN
FindProcess(_In_ ULONG Pid);

EXTERN_C
NTSTATUS
DriverEntry(_In_ PDRIVER_OBJECT DriverObject, _In_ PUNICODE_STRING RegisterPath)
{
	g_Data.m_Lock.Init();

	OB_OPERATION_REGISTRATION Operations[] = {
		{
			PsProcessType,
			OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE,
			OnPreOPenProcess,
			nullptr
		}
	};
	OB_CALLBACK_REGISTRATION Reg = {
		OB_FLT_REGISTRATION_VERSION,
		1,
		RTL_CONSTANT_STRING(L"123456.123"),
		nullptr,
		Operations
	};

	NTSTATUS Status             = STATUS_SUCCESS;

	UNICODE_STRING DeviceName   = RTL_CONSTANT_STRING(L"\\Device\\" PROCESS_PROTECT_NAME);
	UNICODE_STRING SymlinkName  = RTL_CONSTANT_STRING(L"\\??\\" PROCESS_PROTECT_NAME);
	PDEVICE_OBJECT DeviceObject = nullptr;

	do {
		Status = ObRegisterCallbacks(&Reg, &g_Data.m_RegHandle);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to register callbacks (status = %08X)\n", Status));
			break;
		}
		
		Status = IoCreateDevice(DriverObject, 0, &DeviceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to create device object (status = %08X)\n", Status));
			break;
		}

		Status = IoCreateSymbolicLink(&SymlinkName, &DeviceName);
		if (NT_SUCCESS(Status) == FALSE) {
			KdPrint((DRIVER_PREFIX "failed to create symbolic link (status = %08X)\n", Status));
			break;
		}
	} while (FALSE);

	if (NT_SUCCESS(Status) == FALSE) {
		if (DeviceObject != nullptr) {
			IoDeleteDevice(DeviceObject);
		}
		if (g_Data.m_RegHandle != nullptr) {
			ObUnRegisterCallbacks(g_Data.m_RegHandle);
		}
		return Status;
	}

	DriverObject->DriverUnload                         = ProcessProtectUnload;
	DriverObject->MajorFunction[IRP_MJ_CREATE]         = ProcessProtectCreateClose;
	DriverObject->MajorFunction[IRP_MJ_CLOSE]          = ProcessProtectCreateClose;
	DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = ProcessProtectDeviceControl;

	KdPrint((DRIVER_PREFIX "DriverEntry complete successful\n"));

	return Status;
}

VOID
ProcessProtectUnload(_In_ PDRIVER_OBJECT DriverObject)
{
	ObUnRegisterCallbacks(g_Data.m_RegHandle);

	UNICODE_STRING SymlinkName = RTL_CONSTANT_STRING(L"\\??\\" PROCESS_PROTECT_NAME);
	IoDeleteSymbolicLink(&SymlinkName);
	IoDeleteDevice(DriverObject->DeviceObject);

	KdPrint((DRIVER_PREFIX "ProcessProtectUnload complete successful\n"));
}

OB_PREOP_CALLBACK_STATUS
OnPreOPenProcess(_In_ PVOID RegistrationContext, _Inout_ POB_PRE_OPERATION_INFORMATION OperationInformation)
{
	if (OperationInformation->KernelHandle == 1) {
		return OB_PREOP_SUCCESS;
	}

	PEPROCESS Process = (PEPROCESS)OperationInformation->Object;
	ULONG Pid         = HandleToULong(PsGetProcessId(Process));

	AutoLock<FastMutex> Locker(g_Data.m_Lock);
	if (FindProcess(Pid) == TRUE) {
		OperationInformation->Parameters->CreateHandleInformation.DesiredAccess &= ~PROCESS_TERMINATE;
	}

	return OB_PREOP_SUCCESS;
}

NTSTATUS
ProcessProtectDeviceControl(_In_ PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	NTSTATUS Status          = STATUS_SUCCESS;
	ULONG Len                = 0;

	switch (Stack->Parameters.DeviceIoControl.IoControlCode) {
	case IOCTL_PROCESS_PROTECT_BY_PID: {
		ULONG Size = Stack->Parameters.DeviceIoControl.InputBufferLength;
		if (Size % sizeof(ULONG) != 0) {
			Status = STATUS_INVALID_BUFFER_SIZE;
			break;
		}

		ULONG* Data = (ULONG*)Irp->AssociatedIrp.SystemBuffer;

		AutoLock<FastMutex> Locker(g_Data.m_Lock);

		for (ULONG ulI = 0; ulI < Size / sizeof(ULONG); ulI++) {
			ULONG Pid = Data[ulI];
			if (Pid == 0) {
				Status = STATUS_INVALID_PARAMETER;
				break;
			}
			if (FindProcess(Pid) == TRUE) {
				continue;
			}
			if (g_Data.m_PidsCount == g_ulMaxPids) {
				Status = STATUS_TOO_MANY_CONTEXT_IDS;
				break;
			}
			if (AddProcess(Pid) == FALSE) {
				Status = STATUS_UNSUCCESSFUL;
				break;
			}

			Len += sizeof(ULONG);
		}
		break;
	}

	case IOCTL_PROCESS_UNPROTECT_BY_PID: {
		ULONG Size = Stack->Parameters.DeviceIoControl.InputBufferLength;
		if (Size % sizeof(ULONG) != 0) {
			Status = STATUS_INVALID_BUFFER_SIZE;
			break;
		}

		ULONG* Data = (ULONG*)Irp->AssociatedIrp.SystemBuffer;

		AutoLock<FastMutex> Locker(g_Data.m_Lock);

		for (ULONG ulI = 0; ulI < Size / sizeof(ULONG); ulI++) {
			ULONG Pid = Data[ulI];
			if (Pid == 0) {
				Status = STATUS_INVALID_PARAMETER;
				break;
			}
			if (RemoveProcess(Pid) == TRUE) {
				continue;
			}

			Len += sizeof(ULONG);

			if (g_Data.m_PidsCount == 0) {
				break;
			}
		}
		break; 
	}

	case IOCTL_PROCESS_PROTECT_CLEAR: {
		AutoLock<FastMutex> Locker(g_Data.m_Lock);
		::memset(&g_Data.m_Pids, 0, sizeof(g_Data.m_Pids));
		g_Data.m_PidsCount = 0;
		break;
	}

	case IOCTL_PROCESS_QUERY_PIDS: {
		ULONG* Data = (ULONG*)Irp->UserBuffer;
		
		ULONG Pids[g_ulMaxPids] = { 0 };
		for (ULONG ulI = 0; ulI < g_ulMaxPids; ulI++) {
			if (g_Data.m_Pids[ulI] != 0) {
				Pids[ulI] = g_Data.m_Pids[ulI];
			}
		}
		::memcpy(Data, Pids, sizeof(ULONG) * g_ulMaxPids);
		Len = sizeof(ULONG) * g_ulMaxPids;
		break;
	}

	default:
		Status = STATUS_INVALID_DEVICE_REQUEST;
	}

	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = Len;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);

	return Status;
}

NTSTATUS
ProcessProtectCreateClose(_In_ PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	NTSTATUS Status           = STATUS_SUCCESS;
	Irp->IoStatus.Status      = Status;
	Irp->IoStatus.Information = 0;
	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return Status;
}

BOOLEAN
AddProcess(_In_ ULONG Pid)
{
	for (ULONG ulI = 0; ulI < g_ulMaxPids; ulI++) {
		if (g_Data.m_Pids[ulI] == 0) {
			g_Data.m_Pids[ulI] = Pid;
			g_Data.m_PidsCount++;
			return TRUE;
		}
	}
	return FALSE;
}

BOOLEAN
RemoveProcess(_In_ ULONG Pid)
{
	for (ULONG ulI = 0; ulI < g_ulMaxPids; ulI++) {
		if (g_Data.m_Pids[ulI] == Pid) {
			g_Data.m_Pids[ulI] = 0;
			g_Data.m_PidsCount--;
			return TRUE;
		}
	}
	return FALSE;
}

BOOLEAN
FindProcess(_In_ ULONG Pid)
{
	for (ULONG ulI = 0; ulI < g_ulMaxPids; ulI++) {
		if (g_Data.m_Pids[ulI] == Pid) {
			return TRUE;
		}
	}
	return FALSE;
}